import os
os.environ['FLAGS_enable_pir_in_executor'] = '0'
import paddle
import paddle.vision
import paddleslim

paddle.enable_static()

# ----------------- 1. 构建静态图Mobilenet for MNIST -------------------
place = paddle.CUDAPlace()
exe = paddle.static.Executor(place)
startup_prog = paddle.static.Program()
train_prog = paddle.static.Program()

with paddle.static.program_guard(train_prog, startup_prog):
    image = paddle.static.data(name='img', shape=[None, 1, 28, 28], dtype='float32')
    label = paddle.static.data(name='label', shape=[None, 1], dtype='int64')

    # 使用官方mobilenet v1作为特征 backbone
    net = paddle.vision.models.mobilenet_v1(num_classes=10)
    out = net(image)

    loss = paddle.nn.functional.cross_entropy(out, label)
    avg_loss = paddle.mean(loss)

    acc_top1 = paddle.metric.accuracy(out, label, k=1)
    acc_top5 = paddle.metric.accuracy(out, label, k=5)

    fetch_list = [acc_top1, acc_top5, avg_loss]

exe.run(startup_prog)

# ----------------- 2. 剪裁前FLOPs -------------------
flops = paddleslim.analysis.flops(train_prog)
print("FLOPs before prune: ", flops)

# ----------------- 3. 剪裁操作（以全部conv2d都剪去33%为例） -------------------
pruner = paddleslim.prune.Pruner()
# 找所有conv2d权重参数名
params = [param.name for param in train_prog.global_block().all_parameters()
          if 'conv' in param.name and 'weights' in param.name]
ratios = [0.33 for _ in params]

pruned_prog, _, _ = pruner.prune(
    program=train_prog,
    scope=paddle.fluid.global_scope(),
    params=params,
    ratios=ratios,
    place=place
)

# ----------------- 4. 剪裁后FLOPs -------------------
flops = paddleslim.analysis.flops(pruned_prog)
print("FLOPs after prune: ", flops)

# ----------------- 5. 准备MNIST数据并训练一小步 -------------------
import paddle.vision.datasets as datasets
from paddle.io import DataLoader

dataset = datasets.MNIST(mode='train', transform=lambda x: (x/255.).astype('float32'))
loader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=True)

print("Start training (demo 5 batch)...")
for bid, (img_data, label_data) in enumerate(loader()):
    img_np = img_data.numpy()
    img_np = img_np.repeat(3, axis=1)  # [batch, 3, 28, 28]
    label_np = label_data.reshape([-1,1]).numpy()
    acc1, acc5, loss = exe.run(
        pruned_prog,
        feed={'img': img_np, 'label': label_np},
        fetch_list=fetch_list
    )
    print(f"Batch {bid}: acc1={acc1}, acc5={acc5}, loss={loss}")
    if bid>5:
        break